import ast
import re
import random
import string
import argparse
from pathlib import Path

class PythonObfuscator(ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self.var_map = {}
        self.function_params = set()
        self.reserved_names = {
            '__name__', '__main__', 'print', 'input', 
            'open', 'range', 'len', 'str', 'int'
        }
        
    def obfuscate(self, code):
        tree = ast.parse(code)
        self.visit(tree)
        return ast.unparse(tree)
        
    def visit_Name(self, node):
        if (node.id in self.function_params or 
            node.id in dir(__builtins__) or
            node.id in self.reserved_names or
            isinstance(node.ctx, ast.Load) and node.id not in self.var_map):
            return node
            
        if node.id not in self.var_map:
            new_name = ''.join(random.choices(string.ascii_letters, k=8))
            self.var_map[node.id] = new_name
            
        node.id = self.var_map[node.id]
        return node
        
    def visit_FunctionDef(self, node):
        self.function_params = {arg.arg for arg in node.args.args}
        self.generic_visit(node)
        self.function_params = set()
        return node

class CObfuscator:
    def __init__(self):
        self.var_map = {}
        self.func_map = {}
        self.reserved_keywords = {
            'main', 'printf', 'scanf', 'cout', 'cin', 
            'System.out.println', 'public', 'static', 'void'
        }
        
    def obfuscate(self, code):
        # Find all variable declarations
        var_pattern = r'\b(int|float|double|char)\s+([a-zA-Z_]\w*)\b'
        matches = re.finditer(var_pattern, code)
        for match in matches:
            var_type, var_name = match.groups()
            if var_name in self.reserved_keywords:
                continue
            if var_name not in self.var_map:
                new_name = ''.join(random.choices(string.ascii_letters, k=8))
                self.var_map[var_name] = new_name
            code = code.replace(var_name, self.var_map[var_name])
            
        # Find function definitions
        func_pattern = r'\b([a-zA-Z_]\w*)\s*\(([^)]*)\)\s*\{'
        matches = re.finditer(func_pattern, code)
        for match in matches:
            func_name, params = match.groups()
            if func_name in self.reserved_keywords:
                continue
            if func_name not in self.func_map:
                new_name = ''.join(random.choices(string.ascii_letters, k=8))
                self.func_map[func_name] = new_name
            code = code.replace(func_name, self.func_map[func_name])
            
        return code

def obfuscate_file(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            code = f.read()
            
        ext = Path(file_path).suffix.lower()
        if ext == '.py':
            obfuscator = PythonObfuscator()
            obfuscated_code = obfuscator.obfuscate(code)
        elif ext in ('.c', '.cpp', '.java'):
            obfuscator = CObfuscator()
            obfuscated_code = obfuscator.obfuscate(code)
        else:
            print(f"不支持的文件类型: {ext}")
            return False
        
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(obfuscated_code)
            
        print(f"变量名混淆完成: {file_path}")
        return True
        
    except Exception as e:
        print(f"处理文件时出错 {file_path}: {str(e)}")
        return False

def create_test_files():
    """Create test files for all supported languages"""
    test_dir = Path("test")
    test_dir.mkdir(exist_ok=True)
    
    # Python test file
    py_code = """def calculate(a, b):
    result = a + b
    for i in range(3):
        result *= 2
    return result

value = calculate(5, 10)
print(f"计算结果: {value}")
"""
    (test_dir / "python").mkdir(exist_ok=True)
    with open(test_dir / "python" / "sample.py", "w") as f:
        f.write(py_code)
        
    # C test file
    c_code = """#include <stdio.h>

int calculate(int a, int b) {
    int result = a + b;
    for (int i = 0; i < 3; i++) {
        result *= 2;
    }
    return result;
}

int main() {
    int value = calculate(5, 10);
    printf("计算结果: %d\\n", value);
    return 0;
}
"""
    (test_dir / "c").mkdir(exist_ok=True)
    with open(test_dir / "c" / "sample.c", "w") as f:
        f.write(c_code)
        
    # C++ test file
    cpp_code = """#include <iostream>
using namespace std;

class Calculator {
public:
    int calculate(int a, int b) {
        int result = a + b;
        for (int i = 0; i < 3; i++) {
            result *= 2;
        }
        return result;
    }
};

int main() {
    Calculator calc;
    int value = calc.calculate(5, 10);
    cout << "计算结果: " << value << endl;
    return 0;
}
"""
    (test_dir / "cpp").mkdir(exist_ok=True)
    with open(test_dir / "cpp" / "sample.cpp", "w") as f:
        f.write(cpp_code)
        
    # Java test file
    java_code = """public class Sample {
    public static int calculate(int a, int b) {
        int result = a + b;
        for (int i = 0; i < 3; i++) {
            result *= 2;
        }
        return result;
    }

    public static void main(String[] args) {
        int value = calculate(5, 10);
        System.out.println("计算结果: " + value);
    }
}
"""
    (test_dir / "java").mkdir(exist_ok=True)
    with open(test_dir / "java" / "Sample.java", "w") as f:
        f.write(java_code)

from tqdm import tqdm

def process_directory(directory):
    extensions = ('.py', '.c', '.cpp', '.h', '.java')
    files = [p for p in Path(directory).rglob('*') 
             if p.suffix.lower() in extensions]
    
    print(f"找到 {len(files)} 个文件需要处理...")
    with tqdm(files, desc="处理进度") as pbar:
        for path in pbar:
            pbar.set_postfix_str(str(path))
            obfuscate_file(path)

def main():
    parser = argparse.ArgumentParser(
        description='混淆源代码中的变量名(支持Python/C/C++/Java)',
        formatter_class=argparse.RawTextHelpFormatter,
        epilog='''使用示例:
  处理单个文件: python var_obfuscator.py test.py
  处理整个目录: python var_obfuscator.py test/
支持文件类型: .py, .c, .cpp, .h, .java

处理过程中会显示进度条
'''
    )
    parser.add_argument('path', help='要处理的Python文件或目录路径')
    args = parser.parse_args()
    
    path = Path(args.path)
    if path.is_file() and path.suffix == '.py':
        obfuscate_file(path)
    elif path.is_dir():
        process_directory(path)
    else:
        print("错误: 请指定一个有效的.py文件或包含.py文件的目录")

if __name__ == '__main__':
    main()
